package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	"path/filepath"
)

var dir = flag.String("dir", "$$$[wrong]", "Input the target source")

func dirExists(path string) (bool, error) {
	_, err := os.Stat(path)
	if err == nil {
		return true, nil
	}
	if os.IsNotExist(err) {
		err := os.MkdirAll(path, os.ModePerm)
		if err != nil {
			fmt.Printf("mkdir failed![%v]\n", err)
		} else {
			return true, nil
		}
	}
	return false, err
}

func classifyDir(fileInfo *[]os.FileInfo) map[string][]string {
	var fileMap = make(map[string][]string)

	for _, info := range *fileInfo {
		_, fileName := filepath.Split(info.Name())
		ext := filepath.Ext(fileName)
		_, ok := fileMap[ext]
		if ok {
			fileMap[ext] = append(fileMap[ext], fileName)
		} else {
			fileMap[ext] = []string{fileName}
		}
	}

	return fileMap
}

func renameFiles(fileMap *map[string][]string, dir string) {
	index := 0
	fmt.Println("Classify in " + dir + " ...")
	for ext, content := range *fileMap {
		newdir := dir + "/" + ext
		dirExists(newdir)
		for _, file := range content {
			fromPath := dir + "/" + file
			toPath := newdir + "/" + file
			os.Rename(fromPath, toPath)
			fmt.Printf("%v => %v\n", file, ext+"/"+file)
			index++
		}
	}

	fmt.Printf("Success remove %d files\n", index)
}

func readDir(dir string) {
	fileInfo, err := ioutil.ReadDir(dir)

	if err != nil {
		fmt.Printf("Can not found path from %v\n", dir)
		return
	}

	fileMap := classifyDir(&fileInfo)
	renameFiles(&fileMap, dir)

}

func main() {

	flag.Parse()
	if "$$$[wrong]" == *dir {
		fmt.Println("Please input the path by using -dir ...")
		return
	}
	readDir(*dir)

	fmt.Scanln()
}
